import math
from typing import Tuple, Optional
import numpy as np
import osqp
import scipy.sparse as sp  # 仅用于稀疏矩阵转换，不用于矩阵构建


class MPCPathTrackingController:
    def __init__(self, path: np.ndarray, dt: float):
        assert path.shape[0] == 5, "path shape must be (5, N)"
        self.path = path
        self.dt = float(dt)

        # 参数设置
        self.N_c = 10
        self.n_x = 3
        self.n_u = 2
        self.n_var = (self.n_x + self.n_u) * self.N_c
        self.n_con = self.n_var

        # 权重矩阵
        self.Q = np.diag([1.0, 1.0, 2.0])
        self.R = np.diag([0.3, 0.5])

        # 控制约束
        self.lo_bound = np.array([0.2, -0.47])
        self.up_bound = np.array([1.3, 0.47])

        # 构建H矩阵（使用numpy）
        self.H = np.zeros((self.n_var, self.n_var))
        for i in range(self.N_c):
            # 状态代价
            base_x = i * self.n_x
            self.H[base_x:base_x + self.n_x, base_x:base_x + self.n_x] = self.Q
            
            # 控制代价
            base_u = self.n_x * self.N_c + i * self.n_u
            self.H[base_u:base_u + self.n_u, base_u:base_u + self.n_u] = self.R

        self.g = np.zeros(self.n_var)
        
        # OSQP求解器
        self.osqp: Optional[osqp.OSQP] = None
        
        # 统计信息
        self.solve_count = 0
        self.fail_count = 0

    def _setup_osqp_solver(self, A_dense: np.ndarray, lo: np.ndarray, up: np.ndarray) -> bool:
        """正确设置OSQP求解器"""
        try:
            # 将numpy数组转换为scipy稀疏矩阵（仅用于OSQP设置）
            H_sparse = sp.csc_matrix(self.H)
            A_sparse = sp.csc_matrix(A_dense)
            
            if self.osqp is not None:
                # 清理旧的求解器
                try:
                    self.osqp.cleanup()
                except:
                    pass
            
            self.osqp = osqp.OSQP()
            
            # 设置求解器
            self.osqp.setup(
                P=H_sparse, 
                q=self.g, 
                A=A_sparse, 
                l=lo, 
                u=up,
                verbose=False, 
                warm_start=True, 
                max_iter=2000,
                eps_abs=1e-4,
                eps_rel=1e-4,
                polish=True
            )
            
            print(f"OSQP求解器初始化成功，非零元素: {A_sparse.nnz}")
            return True
            
        except Exception as e:
            print(f"OSQP设置失败: {e}")
            return False

    def _update_osqp_constraints(self, A_dense: np.ndarray, lo: np.ndarray, up: np.ndarray) -> bool:
        """更新OSQP约束 - 简化版本：总是重新初始化"""
        try:
            # 由于约束矩阵结构变化较大，每次都重新初始化更稳定
            return self._setup_osqp_solver(A_dense, lo, up)
            
        except Exception as e:
            print(f"OSQP约束更新失败: {e}")
            return False

    def lookup_target_traj(self, t: float, hint_idx: int) -> Tuple[int, np.ndarray, np.ndarray]:
        """查找目标轨迹"""
        time = self.path[0, :]
        x_row = self.path[1, :]
        y_row = self.path[2, :]
        yaw_row = self.path[3, :]
        k_row = self.path[4, :]

        cols = time.shape[0]
        ref_x = np.zeros(3, dtype=float)
        ref_u = np.zeros(2, dtype=float)

        if t <= time[0]:
            ref_x[0] = x_row[0]
            ref_x[1] = y_row[0]
            ref_x[2] = yaw_row[0]
            ref_u[0] = 1.0
            ref_u[1] = math.atan(0.9 * k_row[0])
            return 0, ref_x, ref_u

        if t >= time[cols - 1] or hint_idx >= cols - 1:
            ref_x[0] = x_row[-1]
            ref_x[1] = y_row[-1]
            ref_x[2] = yaw_row[-1]
            ref_u[0] = 0.0
            ref_u[1] = math.atan(0.9 * k_row[-1])
            return cols - 1, ref_x, ref_u

        for i in range(hint_idx, cols - 1):
            if time[i] <= t <= time[i + 1]:
                t_l = time[i]
                t_r = time[i + 1]
                f = (t_r - t) / (t_r - t_l)

                ref_x[0] = x_row[i] * f + x_row[i + 1] * (1.0 - f)
                ref_x[1] = y_row[i] * f + y_row[i + 1] * (1.0 - f)
                ref_x[2] = yaw_row[i] * f + yaw_row[i + 1] * (1.0 - f)

                ref_u[0] = 1.0
                k_interp = 0.9 * (k_row[i] * f + k_row[i + 1] * (1.0 - f))
                ref_u[1] = math.atan(k_interp)
                return i, ref_x, ref_u

        return cols - 1, ref_x, ref_u  # 保底返回

    def _update_matrices(self, t: float, x_c: np.ndarray):
        """更新约束矩阵 - 使用纯numpy"""
        n_x, n_u, N_c, dt = self.n_x, self.n_u, self.N_c, self.dt
        
        # 约束矩阵（稠密）
        A_mat = np.zeros((self.n_con, self.n_var))
        lo = np.zeros(self.n_con)
        up = np.zeros(self.n_con)

        # 第一段参考轨迹
        hint_idx = 0
        idx, ref_x, ref_u = self.lookup_target_traj(t, hint_idx)
        hint_idx = idx

        # 构建线性化矩阵
        A_lin = np.eye(n_x)
        B_lin = np.zeros((n_x, n_u))
        
        A_lin[0, 2] = dt * ref_u[0] * -math.sin(ref_x[2])
        A_lin[1, 2] = dt * ref_u[0] * math.cos(ref_x[2])
        B_lin[0, 0] = dt * math.cos(ref_x[2])
        B_lin[1, 0] = dt * math.sin(ref_x[2])
        B_lin[2, 0] = dt * math.tan(ref_u[1]) / 0.9
        B_lin[2, 1] = dt * ref_u[0] / (0.9 * (math.cos(ref_u[1]) ** 2))

        # 初始状态约束: -x_0 + B u_0 = -A (x_c - ref_x)
        x0_err = x_c - ref_x
        rhs = -A_lin @ x0_err
        
        # 设置初始状态约束
        for j in range(n_x):
            A_mat[j, j] = -1.0  # -x_0
            for k in range(n_u):
                A_mat[j, n_x * N_c + k] = B_lin[j, k]  # +B u_0
            lo[j] = up[j] = rhs[j]

        # 初始控制约束
        lo_u0 = self.lo_bound - ref_u
        up_u0 = self.up_bound - ref_u
        for j in range(n_u):
            idx = n_x * N_c + j
            A_mat[idx, idx] = 1.0
            lo[idx] = lo_u0[j]
            up[idx] = up_u0[j]

        # 后续预测步
        for i in range(1, N_c):
            t_i = t + dt * i
            idx, ref_x, ref_u = self.lookup_target_traj(t_i, hint_idx)
            hint_idx = idx

            # 更新线性化矩阵
            A_lin = np.eye(n_x)
            B_lin = np.zeros((n_x, n_u))
            
            A_lin[0, 2] = dt * ref_u[0] * -math.sin(ref_x[2])
            A_lin[1, 2] = dt * ref_u[0] * math.cos(ref_x[2])
            B_lin[0, 0] = dt * math.cos(ref_x[2])
            B_lin[1, 0] = dt * math.sin(ref_x[2])
            B_lin[2, 0] = dt * math.tan(ref_u[1]) / 0.9
            B_lin[2, 1] = dt * ref_u[0] / (0.9 * (math.cos(ref_u[1]) ** 2))

            # 动力学约束: -x_i + A x_{i-1} + B u_i = 0
            row_base = i * n_x
            for j in range(n_x):
                A_mat[row_base + j, i * n_x + j] = -1.0  # -x_i
                for k in range(n_x):
                    A_mat[row_base + j, (i-1) * n_x + k] = A_lin[j, k]  # +A x_{i-1}
                for k in range(n_u):
                    A_mat[row_base + j, n_x * N_c + i * n_u + k] = B_lin[j, k]  # +B u_i
                lo[row_base + j] = up[row_base + j] = 0.0

            # 控制约束
            lo_u = self.lo_bound - ref_u
            up_u = self.up_bound - ref_u
            for j in range(n_u):
                idx = n_x * N_c + i * n_u + j
                A_mat[idx, idx] = 1.0
                lo[idx] = lo_u[j]
                up[idx] = up_u[j]

        return A_mat, lo, up, ref_x, ref_u

    def step(self, x: float, y: float, yaw: float, t: float) -> Tuple[float, float]:
        """MPC控制步"""
        x_c = np.array([x, y, yaw], dtype=float)
        
        # 获取备用参考控制
        hint_idx = 0
        _, ref_x_backup, ref_u_backup = self.lookup_target_traj(t, hint_idx)
        
        try:
            # 更新约束矩阵
            A_mat, lo, up, ref_x, ref_u = self._update_matrices(t, x_c)
            
            # 打印调试信息
            nnz = np.sum(np.abs(A_mat) > 1e-12)
            if self.solve_count % 10 == 0:
                print(f"约束矩阵非零元素: {nnz}, 时间: {t:.2f}s")
            
            # 更新OSQP约束
            if not self._update_osqp_constraints(A_mat, lo, up):
                print("OSQP设置失败，使用参考控制")
                return float(ref_u_backup[0]), float(ref_u_backup[1])
            
            # 求解QP
            result = self.osqp.solve()
            self.solve_count += 1
            
            if result.info.status_val != osqp.constant("OSQP_SOLVED"):
                self.fail_count += 1
                if self.fail_count % 10 == 0:
                    print(f"QP求解失败: {result.info.status}, 失败次数: {self.fail_count}")
                return float(ref_u_backup[0]), float(ref_u_backup[1])
            
            # 提取控制量
            sol = result.x
            u0 = sol[self.n_x * self.N_c:self.n_x * self.N_c + self.n_u]
            v_cmd = ref_u[0] + u0[0]
            steer_cmd = ref_u[1] + u0[1]
            
            # 安全约束
            v_cmd = np.clip(v_cmd, self.lo_bound[0], self.up_bound[0])
            steer_cmd = np.clip(steer_cmd, self.lo_bound[1], self.up_bound[1])
            
            return float(v_cmd), float(steer_cmd)
            
        except Exception as e:
            print(f"MPC步执行异常: {e}")
            return float(ref_u_backup[0]), float(ref_u_backup[1])

    def get_stats(self) -> dict:
        """获取控制器统计信息"""
        total = max(self.solve_count, 1)
        success_rate = (total - self.fail_count) / total
        return {
            "total_solves": self.solve_count,
            "failures": self.fail_count,
            "success_rate": success_rate
        }

    def cleanup(self):
        """清理资源"""
        if self.osqp is not None:
            try:
                self.osqp.cleanup()
            except:
                pass
